import numpy as np

from htm.Connection import Connection


class ApicalTMPairRegion:

    def __init__(self, params):
        # Parameters
        self.columnCount = params['columnCount']
        self.cellsPerColumn = params['cellsPerColumn']
        self.reducedBasalThreshold = params['reducedBasalThreshold']
        # Initialization
        self.basalConnection = Connection(params['basalConnParams'])
        self.apicalConnection = Connection(params['apicalConnParams'])
        self.learn = True
        self.useApicalTiebreak = True
        self.useApicalModulationBasalThreshold = True
        # Mutable States
        self.predictedCells = np.array([])
        self.activeBasalSegments = np.array([])
        self.activeApicalSegments = np.array([])
        self.matchingBasalSegments = np.array([])
        self.matchingApicalSegments = np.array([])
        self.basalPotentialOverlaps = None
        self.apicalPotentialOverlaps = None
        self.activeCells = np.array([])
        self.winnerCells = np.array([])
        self.predictedActiveCells = np.array([])

    def compute(self, feedforward, basalInput, apicalInput, basalGrowthCandidates, apicalGrowthCandidates):
        if basalGrowthCandidates is None:
            basalGrowthCandidates = basalInput
        if apicalGrowthCandidates is None:
            apicalGrowthCandidates = apicalInput
        self.depolarizeCells(basalInput, apicalInput)
        self.activateCells(feedforward, basalInput, apicalInput, basalGrowthCandidates, apicalGrowthCandidates)

    def depolarizeCells(self, basalInput, apicalInput):
        # Apical Segments Calculation
        apicalConnectedOverlaps, apicalPotentialOverlaps = self.apicalConnection.computeActivity(apicalInput)
        activeApicalSegments = Connection.filterOverlapByThreshold(apicalConnectedOverlaps, self.apicalConnection.activationThreshold)
        matchingApicalSegments = Connection.filterOverlapByThreshold(apicalPotentialOverlaps, self.apicalConnection.minThreshold)
        if self.learn or not self.useApicalModulationBasalThreshold:
            reducedBasalThresholdCells = ()
        else:
            reducedBasalThresholdCells = self.apicalConnection.mapSegmentsToCells(activeApicalSegments)
        # Basal Segments Calculation
        basalConnectedOverlaps, basalPotentialOverlaps = self.basalConnection.computeActivity(basalInput)
        outrightBasalActiveSegments = Connection.filterOverlapByThreshold(basalConnectedOverlaps, self.basalConnection.activationThreshold)
        if self.reducedBasalThreshold != self.basalConnection.activationThreshold and len(reducedBasalThresholdCells) > 0:
            potentiallyActiveSegments = Connection.filterOverlapByInterval(basalConnectedOverlaps, self.reducedBasalThreshold, self.basalConnection.activationThreshold)
            conditionallyActiveSegments = self.getConditActiveBasalSegments(potentiallyActiveSegments, reducedBasalThresholdCells)
            activeBasalSegments = np.union1d(outrightBasalActiveSegments, conditionallyActiveSegments)
        else:
            activeBasalSegments = outrightBasalActiveSegments
        matchingBasalSegments = Connection.filterOverlapByThreshold(basalPotentialOverlaps, self.basalConnection.minThreshold)
        # Predicted Cells Calculation
        cellsForActiveBasalSegments = self.basalConnection.mapSegmentsToCells(activeBasalSegments)
        cellsForActiveApicalSegments = self.apicalConnection.mapSegmentsToCells(activeApicalSegments)
        fullyDepolarizedCells = np.intersect1d(cellsForActiveBasalSegments, cellsForActiveApicalSegments)
        partlyDepolarizedCells = np.setdiff1d(cellsForActiveBasalSegments, fullyDepolarizedCells)
        cellsWithNoInhibition = self.basalConnection.getCellsWithNoInhibition(partlyDepolarizedCells, fullyDepolarizedCells)
        self.predictedCells = np.union1d(fullyDepolarizedCells, cellsWithNoInhibition)
        if not self.useApicalTiebreak:
            self.predictedCells = cellsForActiveBasalSegments
        self.activeBasalSegments = activeBasalSegments
        self.activeApicalSegments = activeApicalSegments
        self.matchingBasalSegments = matchingBasalSegments
        self.matchingApicalSegments = matchingApicalSegments
        self.basalPotentialOverlaps = basalPotentialOverlaps
        self.apicalPotentialOverlaps = apicalPotentialOverlaps

    def activateCells(self, feedforward, basalInput, apicalInput, basalGrowthCandidates, apicalGrowthCandidates):
        # Calculate active cells
        correctPredictedCells = self.predictedCells[np.in1d(self.predictedCells // self.cellsPerColumn, feedforward)]
        burstingColumns = np.setdiff1d(feedforward, self.predictedCells // self.cellsPerColumn)
        cellsInBurstingColumns = Connection.getCellsInColumns(burstingColumns, self.cellsPerColumn)
        newActiveCells = np.union1d(correctPredictedCells, cellsInBurstingColumns)
        # Calculate basal learning
        learningActiveBasalSegments = self.basalConnection.filterSegmentsByCell(self.activeBasalSegments, correctPredictedCells)
        cellsForMatchingBasalSegments = self.basalConnection.mapSegmentsToCells(self.matchingBasalSegments)
        matchingCells = np.unique(cellsForMatchingBasalSegments)
        matchingCellsInBurstingColumns = matchingCells[np.in1d(matchingCells // self.cellsPerColumn, burstingColumns)]
        learningMatchingBasalSegments = self.basalConnection.chooseBestSegmentPerColumn(matchingCellsInBurstingColumns, self.matchingBasalSegments, self.basalPotentialOverlaps)
        learningMatchingCells = self.basalConnection.mapSegmentsToCells(learningMatchingBasalSegments)
        burstingColumnsWithNoMatch = np.setdiff1d(burstingColumns, matchingCells // self.cellsPerColumn)
        newBasalSegmentCells = self.basalConnection.getCellsWithFewestSegments(burstingColumnsWithNoMatch)
        learningCells = np.union1d(np.union1d(correctPredictedCells, learningMatchingCells), newBasalSegmentCells)
        correctMatchingBasalMask = np.in1d(cellsForMatchingBasalSegments // self.cellsPerColumn, feedforward)
        basalSegmentsToPunish = self.matchingBasalSegments[~correctMatchingBasalMask]
        # Calculate apical learning
        learningActiveApicalSegments = self.apicalConnection.filterSegmentsByCell(self.activeApicalSegments, learningCells)
        learningActiveApicalCells = self.apicalConnection.mapSegmentsToCells(learningActiveApicalSegments)
        learningCellsWithoutActiveApical = np.setdiff1d(learningCells, learningActiveApicalCells)
        cellsForMatchingApicalSegments = self.apicalConnection.mapSegmentsToCells(self.matchingApicalSegments)
        learningCellsWithMatchingApical = np.intersect1d(learningCellsWithoutActiveApical, cellsForMatchingApicalSegments)
        learningMatchingApicalSegments = self.apicalConnection.chooseBestSegmentPerColumn(learningCellsWithMatchingApical, self.matchingApicalSegments, self.apicalPotentialOverlaps)
        newApicalSegmentCells = np.setdiff1d(learningCellsWithoutActiveApical, learningCellsWithMatchingApical)
        correctMatchingApicalMask = np.in1d(cellsForMatchingApicalSegments // self.cellsPerColumn, feedforward)
        apicalSegmentsToPunish = self.matchingApicalSegments[~correctMatchingApicalMask]
        # Learn
        if self.learn:
            self.basalConnection.learnOnExisting(learningActiveBasalSegments, basalInput, basalGrowthCandidates)
            self.basalConnection.learnOnExisting(learningMatchingBasalSegments, basalInput, basalGrowthCandidates)
            self.apicalConnection.learnOnExisting(learningActiveApicalSegments, apicalInput, apicalGrowthCandidates)
            self.apicalConnection.learnOnExisting(learningMatchingApicalSegments, apicalInput, apicalGrowthCandidates)
            if self.basalConnection.permanencePunishment != 0.0:
                self.basalConnection.adjustActiveSynapses(basalSegmentsToPunish, basalInput, -self.basalConnection.permanencePunishment, 0)
            if self.apicalConnection.permanencePunishment != 0.0:
                self.apicalConnection.adjustActiveSynapses(apicalSegmentsToPunish, apicalInput, -self.apicalConnection.permanencePunishment, 0)
            self.basalConnection.learnOnNew(newBasalSegmentCells, basalGrowthCandidates)
            self.apicalConnection.learnOnNew(newApicalSegmentCells, apicalGrowthCandidates)
        # For Output
        self.activeCells = newActiveCells
        self.winnerCells = learningCells
        self.predictedActiveCells = correctPredictedCells

    def getConditActiveBasalSegments(self, potentiallyActiveSegments, reducedBasalThresholdCells):
        res = np.array([])
        for segId in potentiallyActiveSegments:
            cellId = self.basalConnection.getCellIdBySegmentId(segId)
            if cellId in reducedBasalThresholdCells:
                res = np.append(res, segId)
        return res

    def reset(self):
        self.predictedCells = np.array([])
        self.activeBasalSegments = np.array([])
        self.activeApicalSegments = np.array([])
        self.matchingBasalSegments = np.array([])
        self.matchingApicalSegments = np.array([])
        self.basalPotentialOverlaps = None
        self.apicalPotentialOverlaps = None
        self.activeCells = np.array([])
        self.winnerCells = np.array([])
        self.predictedActiveCells = np.array([])

    def getBasalPredictedCells(self):
        return np.unique(self.basalConnection.mapSegmentsToCells(self.activeBasalSegments))
